package main

import (
	"bufio"
	"context"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"strings"
	"sync"
	"sync/atomic"
)

const (
	headerReadLen = 512 // MIME类型检测长度
	maxOpenFiles  = 100 // 最大并行打开文件数
)

type (
	task struct {
		index int32
		path  string
	}

	result struct {
		index   int32
		content chan []byte
		err     error
	}
)

func init() {
	// 重写帮助信息
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "代码合并工具 - 将指定目录的代码文件合并为Markdown文档\n\n")
		fmt.Fprintf(os.Stderr, "使用方法: %s [选项] 目标目录\n\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "选项:\n")
		flag.PrintDefaults()
		fmt.Fprintf(os.Stderr, "\n示例:\n")
		fmt.Fprintf(os.Stderr, "  # 处理整个目录\n")
		fmt.Fprintf(os.Stderr, "  %s -mergtype go,js ~/project\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  # 指定输出路径\n")
		fmt.Fprintf(os.Stderr, "  %s -o docs/output.md ~/src\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  # 处理所有文本文件\n")
		fmt.Fprintf(os.Stderr, "  %s -mergtype all ~/project\n\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  # 匹配src目录下的go文件\n")
		fmt.Fprintf(os.Stderr, "  %s -regex '^src/.*\\.go$' ~/project\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  -exclude string\n    排除目录模式（逗号分隔，如 'node_modules,test*'）\n")
		fmt.Fprintf(os.Stderr, "  # 排除node_modules目录\n")
		fmt.Fprintf(os.Stderr, "  %s -exclude node_modules ~/project\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "  # 排除所有test开头的目录\n")
		fmt.Fprintf(os.Stderr, "  %s -exclude 'test*' ~/src\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "注意事项:\n")
		fmt.Fprintf(os.Stderr, "  - 文件类型过滤基于扩展名（如.go .js）\n")
		fmt.Fprintf(os.Stderr, "  - 不指定-mergtype时处理所有文本文件\n")
	}
}

func main() {
	var (
		outputPath     string
		mergeTypesStr  string
		showHelp       bool
		regexStr       string
		excludeDirsStr string
	)
	flag.StringVar(&regexStr, "regex", "", "正则表达式匹配文件相对路径")
	flag.StringVar(&excludeDirsStr, "exclude", "node_module,venv", "排除目录列表（逗号分隔，支持通配符模式）")
	flag.StringVar(&outputPath, "o", "mergedown.md", "输出文件路径")
	flag.StringVar(&mergeTypesStr, "mergtype", "", "要合并的文件类型（逗号分隔）")
	flag.BoolVar(&showHelp, "help", false, "显示帮助信息")
	flag.BoolVar(&showHelp, "h", false, "显示帮助信息（简写）")
	flag.Parse()

	// 显示帮助信息
	if showHelp {
		flag.Usage()
		return
	}

	// 参数校验与初始化
	args := flag.Args()
	if len(args) != 1 {
		flag.Usage()
		log.Fatal("必须且只能指定一个目标目录")
	}
	targetDir := filepath.Clean(args[0])
	err := validateDirectory(targetDir)
	if err != nil {
		log.Fatal(err)
	}

	// 处理合并类型 - 修复点：特殊处理"all"关键字
	mergeTypes := make(map[string]struct{})
	mergeAll := false
	for _, t := range strings.Split(mergeTypesStr, ",") {
		t = strings.ToLower(strings.TrimSpace(t))
		if t == "" {
			continue
		}
		if t == "all" {
			mergeAll = true
			// 一旦出现"all"，清空之前的类型并跳出循环
			mergeTypes = make(map[string]struct{})
			break
		}
		mergeTypes[t] = struct{}{}
	}

	// 解析排除参数（main函数参数校验部分后）
	var excludePatterns []string
	if excludeDirsStr != "" {
		excludePatterns = strings.Split(excludeDirsStr, ",")
		for i := range excludePatterns {
			excludePatterns[i] = strings.TrimSpace(excludePatterns[i])
			// 转换为小写匹配（可选）
			excludePatterns[i] = strings.ToLower(excludePatterns[i])
		}
	}
	var re *regexp.Regexp
	if regexStr != "" {
		var err error
		if re, err = regexp.Compile(regexStr); err != nil {
			log.Fatalf("无效的正则表达式: %v", err)
		}
	}

	// 设置默认输出路径
	if outputPath == "." || outputPath == "" {
		// 获取父目录名称
		parentDir := filepath.Dir(targetDir)
		parentBase := filepath.Base(parentDir)

		// 处理特殊情况
		switch {
		case parentBase == ".": // 如目标目录是相对路径 "./project"
			parentBase = filepath.Base(targetDir)
		case parentBase == "": // 处理根目录情况
			parentBase = filepath.Base(targetDir)
			if parentBase == "" {
				log.Fatal("无法生成输出文件名：目标目录路径无效")
			}
		}
		// 构建最终路径
		outputPath = filepath.Join(".", parentBase+".md")
		outputPath = filepath.Clean(outputPath) // 再次规范化路径

		fmt.Printf("[DEBUG] 特殊路径处理结果：%s\n", outputPath) // 调试用，完成后可删除
	}

	outputPath = filepath.Clean(outputPath)

	// 创建输出目录
	if err := os.MkdirAll(filepath.Dir(outputPath), 0755); err != nil {
		log.Fatalf("创建输出目录失败: %v", err)
	}

	// 创建输出文件
	outFile, err := os.Create(outputPath)
	if err != nil {
		log.Fatal("创建输出文件失败:", err)
	}
	defer outFile.Close()

	// 上下文管理
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// 创建管道
	taskCh := make(chan task, maxOpenFiles)
	resultCh := make(chan result, maxOpenFiles)

	// 启动收集器
	go collector(ctx, outFile, resultCh)

	// 启动处理器池
	var wg sync.WaitGroup
	for i := 0; i < runtime.NumCPU(); i++ {
		wg.Add(1)
		go processor(ctx, &wg, targetDir, mergeAll, mergeTypes, taskCh, resultCh)
	}

	// 遍历目录生成任务
	producer(ctx, targetDir, mergeAll, mergeTypes, re, excludePatterns, taskCh)

	// 等待处理完成
	wg.Wait()
	close(resultCh)

	fmt.Printf("成功生成汇总文件: %s\n", outputPath)
}

// 新增目录验证函数
func validateDirectory(path string) error {
	info, err := os.Stat(path)
	if err != nil {
		if os.IsNotExist(err) {
			return fmt.Errorf("目录不存在: %q", path)
		}
		return fmt.Errorf("无法访问目录: %q (%v)", path, err)
	}

	if !info.IsDir() {
		return fmt.Errorf("路径不是目录: %q", path)
	}

	// 检查可读权限
	f, err := os.Open(path)
	if err != nil {
		return fmt.Errorf("目录不可读: %q (%v)", path, err)
	}
	f.Close()

	return nil
}

// 4. 新增路径匹配函数
func isExcluded(relPath string, patterns []string) bool {
	normalized := filepath.ToSlash(strings.ToLower(relPath))

	for _, pattern := range patterns {
		matched, _ := filepath.Match(
			strings.ToLower(pattern),
			normalized,
		)
		if matched {
			return true
		}

		// 支持目录通配（如排除所有子目录）
		if strings.HasPrefix(pattern, "*/") {
			if matched, _ := filepath.Match("*"+pattern, normalized); matched {
				return true
			}
		}
	}
	return false
}

// 生产者：遍历目录生成任务
func producer(ctx context.Context, targetDir string, mergeAll bool, mergeTypes map[string]struct{}, re *regexp.Regexp, excludePatterns []string, taskCh chan<- task) {
	defer close(taskCh)
	var index int32

	filepath.Walk(targetDir, func(path string, info os.FileInfo, err error) error {
		// 获取相对路径（用于匹配）
		relPath, err := filepath.Rel(targetDir, path)
		if err != nil {
			return nil
		}
		// 新增排除逻辑：目录匹配
		if info.IsDir() {
			if isExcluded(relPath, excludePatterns) {
				return filepath.SkipDir
			}
			return nil
		}
		// 文件过滤（原有逻辑+排除校验）
		if isExcluded(relPath, excludePatterns) || !validFileType(relPath, mergeAll, mergeTypes, re) {
			return nil
		}

		select {
		case taskCh <- task{atomic.AddInt32(&index, 1), path}:
		case <-ctx.Done():
			return ctx.Err()
		}
		return nil
	})
}

// 处理器：并发处理文件
func processor(ctx context.Context, wg *sync.WaitGroup, targetDir string, mergeAll bool, mergeTypes map[string]struct{},
	taskCh <-chan task, resultCh chan<- result) {

	defer wg.Done()

	for {
		select {
		case t, ok := <-taskCh:
			if !ok {
				return
			}

			relPath, err := filepath.Rel(targetDir, t.path)
			if err != nil {
				sendResult(ctx, resultCh, result{t.index, nil, err})
				continue
			}

			contentCh := make(chan []byte, 4)
			sendResult(ctx, resultCh, result{t.index, contentCh, nil})

			if err := processFile(ctx, t.path, relPath, contentCh); err != nil {
				sendResult(ctx, resultCh, result{t.index, nil, err})
			}
			close(contentCh)

		case <-ctx.Done():
			return
		}
	}
}

// 收集器：按顺序写入文件
func collector(ctx context.Context, outFile *os.File, resultCh <-chan result) {
	var (
		current int32 = 1
		cache         = make(map[int32]chan []byte)
		writer        = bufio.NewWriterSize(outFile, 64*1024)
	)

	defer writer.Flush()

	for {
		select {
		case res, ok := <-resultCh:
			if !ok {
				return
			}

			if res.err != nil {
				log.Fatal("处理错误:", res.err)
			}

			cache[res.index] = res.content

			for {
				ch, exists := cache[current]
				if !exists {
					break
				}

				for data := range ch {
					if _, err := writer.Write(data); err != nil {
						log.Fatal("写入错误:", err)
					}
				}

				delete(cache, current)
				current++
			}

		case <-ctx.Done():
			return
		}
	}
}

// 流式处理文件
func processFile(ctx context.Context, path, relPath string, contentCh chan<- []byte) error {
	file, err := os.Open(path)
	if err != nil {
		return fmt.Errorf("打开文件失败: %v", err)
	}
	defer file.Close()

	// 写入文件头
	header := fmt.Sprintf("## %s\n\n", filepath.ToSlash(relPath))
	select {
	case contentCh <- []byte(header):
	case <-ctx.Done():
		return ctx.Err()
	}

	// 检测二进制文件
	isText, err := isTextFile(file)
	if err != nil {
		return err
	}
	if !isText {
		select {
		case contentCh <- []byte("```\n[二进制文件内容已忽略]\n```\n\n"):
		case <-ctx.Done():
		}
		return nil
	}

	// 写入代码块开始
	ext := strings.TrimPrefix(filepath.Ext(path), ".")
	codeStart := "```" + ext + "\n"
	select {
	case contentCh <- []byte(codeStart):
	case <-ctx.Done():
		return ctx.Err()
	}

	// 流式复制内容
	reader := bufio.NewReader(file)
	buf := make([]byte, 4096)
	for {
		n, err := reader.Read(buf)
		if n > 0 {
			data := make([]byte, n)
			copy(data, buf[:n])
			select {
			case contentCh <- data:
			case <-ctx.Done():
				return ctx.Err()
			}
		}
		if err != nil {
			if err == io.EOF {
				break
			}
			return fmt.Errorf("读取失败: %v", err)
		}
	}

	// 处理结尾换行
	if _, err := file.Seek(-1, io.SeekEnd); err == nil {
		var last [1]byte
		if _, err := file.Read(last[:]); err == nil && last[0] != '\n' {
			select {
			case contentCh <- []byte("\n"):
			case <-ctx.Done():
			}
		}
	}

	// 写入代码块结束
	select {
	case contentCh <- []byte("```\n\n"):
	case <-ctx.Done():
	}
	return nil
}

// 辅助函数
func sendResult(ctx context.Context, ch chan<- result, res result) {
	select {
	case ch <- res:
	case <-ctx.Done():
	}
}

// 修复点：修改validFileType函数 - 支持all关键字和空类型处理
func validFileType(relPath string, mergeAll bool, mergeTypes map[string]struct{}, re *regexp.Regexp) bool {
	// 4. 正则匹配（如果存在）
	if re != nil && !re.MatchString(relPath) {
		return false
	}

	// 1. 如果是all模式，跳过扩展名检查
	// 2. 如果mergeTypes为空（未指定类型），也视为处理所有文件
	if mergeAll || len(mergeTypes) == 0 {
		return true
	}

	// 3. 普通模式：检查扩展名
	ext := strings.ToLower(strings.TrimPrefix(filepath.Ext(relPath), "."))
	if _, exists := mergeTypes[ext]; !exists {
		return false
	}

	return true
}

func isTextFile(file *os.File) (bool, error) {
	header := make([]byte, headerReadLen)
	n, err := file.Read(header)
	if err != nil && err != io.EOF {
		return false, err
	}
	if _, err := file.Seek(0, io.SeekStart); err != nil {
		return false, err
	}
	isText := strings.HasPrefix(http.DetectContentType(header[:n]), "text/")
	return isText, nil
}
